#!/usr/bin/env python3.7
# -*- coding:UTF=8 -*-
"""
Author: YangYunYi
Date: 2022/5/12 9:55
docs: draw.py.py
description: 
"""

import xlrd
from os.path import abspath
from pyecharts import options as opts
from pyecharts.charts import TreeMap
from pyecharts.commons.utils import JsCode

percent_jscode_formatter = JsCode("function(x){return x.name+'('+x.value+'%)';}")
value_jscode_formatter = JsCode("function(x){return x.name+'('+x.value+')';}")


def open_file(path, sheet_name):
    # Open the workbook
    xl_workbook = xlrd.open_workbook(path)

    # List sheet names, and pull a sheet by name
    #
    sheet_names = xl_workbook.sheet_names()
    xl_sheet = xl_workbook.sheet_by_name(sheet_name)
    print('Sheet Names', sheet_names)
    return xl_sheet


def get_data(file_path):
    xl_sheet = open_file(file_path, "data")

    data_tree = {}
    num_cols = xl_sheet.ncols  # Number of columns
    type_stack = [""] * num_cols
    for row_idx in range(0, xl_sheet.nrows):  # Iterate through rows
        for col_idx in range(0, num_cols - 1):  # Iterate through columns

            cell_obj = xl_sheet.cell(row_idx, col_idx).value  # Get cell object by row, col
            if cell_obj != "":
                type_stack[col_idx] = cell_obj
            else:
                continue

            if col_idx == 0:
                last_dict = data_tree
            else:
                last_dict = data_tree
                for i in range(0, col_idx):
                    last_dict = last_dict[type_stack[i]]
            if col_idx == num_cols - 2:
                if cell_obj not in last_dict:
                    last_dict[cell_obj] = xl_sheet.cell(row_idx, col_idx + 1).value
            else:
                if cell_obj not in last_dict:
                    last_dict[cell_obj] = {}

    print(data_tree)
    ret = traverse_dict(data_tree)
    return ret


def get_pic_size(size_str):
    size_list = size_str.split("*")
    if len(size_list) != 2:
        return 1920, 1080
    return int(size_list[0]), int(size_list[1])


def get_config(file_path):
    xl_sheet = open_file(file_path, "config")
    config_dict = {"pic_title": "demo_title",
                   "sub_title": "demo_sub",
                   "label_type": "none",
                   "pic_size": {"width": 1920, "height": 1080}
                   }
    for row_idx in range(1, xl_sheet.nrows):  # Iterate through rows
        config_key = xl_sheet.cell(row_idx, 0).value
        config_value = xl_sheet.cell(row_idx, 1).value
        if config_key not in config_dict.keys():
            print("Invalid config item.")
            continue
        if config_key == "pic_size":
            w, h = get_pic_size(config_value)
            config_dict[config_key] = {"width": "{}px".format(w), "height": "{}px".format(h)}
        else:
            config_dict[config_key] = config_value

    return config_dict


def traverse_dict(dict_data):
    ret = []
    for key, value in dict_data.items():
        node = {"name": key}
        if isinstance(value, dict):
            children = traverse_dict(value)
            total_value = 0
            for child in children:
                total_value += child["value"]
            node["value"] = total_value
            node["children"] = children
        else:
            node["value"] = value
        ret.append(node)
    return ret


def get_total(data_list):
    total = 0
    for data_item in data_list:
        total += data_item["value"]
    return total


def trans_to_percent(data, total):
    value = data["value"]
    if "children" in data.keys():
        for child in data["children"]:
            # child_value = child["value"]
            trans_to_percent(child, value)
    data["value"] = "%.2f" % (value / total * 100)


def get_percent(data):
    total = get_total(data)
    for item in data:
        trans_to_percent(item, total)


def draw_pic(data, config):
    formatter = None
    if config["label_type"] == "value":
        formatter = value_jscode_formatter
    elif config["label_type"] == "percent":
        get_percent(data)
        formatter = percent_jscode_formatter
    print(data)
    print(config)
    c = (
        TreeMap(init_opts=opts.InitOpts(width=config["pic_size"]["width"], height=config["pic_size"]["height"]))
            .add(
            series_name=config["pic_title"],
            data=data,
            label_opts=opts.LabelOpts(position="inside", formatter=formatter),
            levels=[
                opts.TreeMapLevelsOpts(
                    treemap_itemstyle_opts=opts.TreeMapItemStyleOpts(
                        border_color="#555", border_width=4, gap_width=4
                    )
                ),
                opts.TreeMapLevelsOpts(
                    color_saturation=[0.3, 0.6],
                    treemap_itemstyle_opts=opts.TreeMapItemStyleOpts(
                        border_color_saturation=0.7, gap_width=2, border_width=2
                    ),
                ),
                opts.TreeMapLevelsOpts(
                    color_saturation=[0.3, 0.5],
                    treemap_itemstyle_opts=opts.TreeMapItemStyleOpts(
                        border_color_saturation=0.6, gap_width=1
                    ),
                ),
                opts.TreeMapLevelsOpts(color_saturation=[0.3, 0.5]),
            ],
        )
            .set_global_opts(title_opts=opts.TitleOpts(title=config["sub_title"]))
            .render("treemap_output.html")
    )


def main():
    f_path = abspath("file.xls")
    data = get_data(f_path)
    config = get_config(f_path)
    draw_pic(data, config)


if __name__ == "__main__":
    main()
